package com.trendrr.beanstalk;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.NotYetConnectedException;
import java.nio.channels.SocketChannel;
import java.util.logging.Level;
import java.util.logging.Logger;
/**
* Wraps the beanstalk connection.
*
* @author dustin
*/
public class BeanstalkConnection {
protected Logger log = Logger.getLogger("com.trendrr.beanstalk");
private SocketChannel channel;
private ByteArrayOutputStream outbuf = new ByteArrayOutputStream();
private void throwException(Exception x) throws BeanstalkException {
if (x instanceof NotYetConnectedException) throw new BeanstalkDisconnectedException(x);
if (x instanceof IOException) throw new BeanstalkDisconnectedException(x);
throw new BeanstalkException(x);
}
public void close() {
try {
outbuf.close();
} catch (Exception x) {
log.log(Level.FINER, "Caught", x);
}
try {
channel.close();
} catch (Exception x) {
log.log(Level.FINER, "Caught", x);
}
}
public void connect(String addr, int port) throws BeanstalkException {
try {
this.channel = SocketChannel.open();
this.channel.connect(new InetSocketAddress(addr, port));
this.channel.finishConnect();
} catch (Exception x) {
throw new BeanstalkException(x);
}
}
@Override
public void finalize() throws Throwable {
super.finalize();
this.close();
}
public boolean isOpen() {
return channel != null && channel.isOpen();
}
public byte[] readBytes(int numBytes) throws BeanstalkException {
byte[] bytes = new byte[numBytes];
byte[] array = this.outbuf.toByteArray();
this.outbuf = new ByteArrayOutputStream();
int bytesWritten = 0;
for (int i = 0; i < array.length; i++) {
if (bytesWritten < bytes.length) {
bytes[i] = array[i];
bytesWritten++;
} else this.outbuf.write(array[i]);
}
if (bytesWritten >= bytes.length) return bytes;
int numRead = 1;
while (numRead > 0) {
ByteBuffer buf = ByteBuffer.allocate(4096);
try {
numRead = channel.read(buf);
} catch (Exception x) {
this.throwException(x);
}
byte[] read = buf.array();
for (int i = 0; i < numRead; i++) {
if (bytesWritten < bytes.length) bytes[bytesWritten] = read[i];
else this.outbuf.write(read[i]);
bytesWritten++;
}
if (bytesWritten >= bytes.length) {
log.finer("468 GOT : " + bytesWritten + " " + bytes.length);
return bytes;
}
}
return bytes;
}
/**
* returns the control response. ends with \r\n
*
* @return
*/
public String readControlResponse() throws BeanstalkException {
String response = null;
int count = 0;
while (response == null) {
count++;
outbuf = new ByteArrayOutputStream();
ByteBuffer buf = ByteBuffer.allocate(4096);
if (count > 10000) {
throw new BeanstalkException("Buffer has been empty for more than 100 seconds.");
}
try {
if (channel.read(buf) == 0) {
log.warning("Nothing in the buffer, sleeping for 100 millis; will try again.");
try {
Thread.sleep(100);
} catch (Exception x) {
log.log(Level.SEVERE, "Caught", x);
}
continue;
}
} catch (Exception x) {
this.throwException(x);
}
byte[] bytes = buf.array();
ByteArrayOutputStream stringBuf = new ByteArrayOutputStream();
byte lastByte = ' ';
for (int i = 0; i < buf.position(); i++) {
byte curByte = bytes[i];
if (lastByte == '\r' && curByte == '\n' && response == null) {
response = new String(stringBuf.toByteArray()).trim();
if (response.isEmpty()) {
log.warning("Errant line end found, possibly from the previous request. Skipping.");
response = null;
}
continue;
}
if (response == null) stringBuf.write(curByte);
else outbuf.write(curByte);
lastByte = curByte;
}
}
return response;
}
public void write(byte[] bytes) throws BeanstalkException {
try {
ByteBuffer buf = ByteBuffer.wrap(bytes);
while (buf.hasRemaining()) channel.write(buf);
} catch (Exception x) {
this.throwException(x);
}
}
public void write(String str) throws BeanstalkException {
try {
ByteBuffer buf = ByteBuffer.wrap(str.getBytes());
while (buf.hasRemaining()) channel.write(buf);
} catch (Exception x) {
this.throwException(x);
}
}
}